'''
Created on 07.09.2012

@author: Julia
'''
from __future__ import division
import os
import numpy as np
import scipy.constants as sconst
import Physics as phy

d=1


class SpinType():
    '''
    saves and loads spin types like NV etc. with all the important parameters
    '''
    def __init__(self):
        '''
        Constructor
            name: String with name of spin
            spinNumber: S or I
            zfs: zero-field-splitting
        '''
        self.name="spin"
        self.spinNumber=0.5
        self.gyro=1
        self.zfs=0

    def get_savetext(self):
        text=self.name+"\t"+str(self.spinNumber)+"\t"+str(self.gyro)+"\t"+str(self.zfs)+"\n"
        return text
    
    def save_art(self, filename="libs/Spins.txt"):
        text=self.get_savetext()
        if os.path.exists(filename):
            f=open(filename, "r")
            spin_list=f.read()
            f.close()
            spin_new=spin_list+text
        else:
            if not os.path.exists(os.path.dirname(filename)):
                os.makedirs(os.path.dirname(filename))
            spin_new=text
        f=open(filename, "w")
        f.write(spin_new)
        f.close()

    def read_savetext(self, text):
        spin_info=text.split("\t")
        if len(spin_info)==4:
            self.name=spin_info[0]
            self.spinNumber=float(spin_info[1])
            self.gyro=float(spin_info[2])
            self.zfs=float(spin_info[3])
        else:
            print "Spin has not the right format"
    
    def load_art(self, i, filename="libs/Spins.txt"):
        try:
            f=open(filename, "r")
            spin_list=f.readlines()
            f.close()
            self.read_savetext(spin_list[i])
        except:
            print "Cannot load spin"
        
        
class OneSpin():
    """
    classdocs
    """        
    def __init__(self, art, sys=False):
        """
        Constructor
        """
        self.art=art
        self.position=(0, 0, 0)
        self.system = sys
        self.active=True
        self.orientation_axis=(0, 0, 1)
        self.x, self.y, self.z=self.position
        self.get_r()
        self.state=np.zeros((self.art.spinNumber*2+1, self.art.spinNumber*2+1), dtype=complex)
        self.state[0, 0]=1
        #self.r=np.sqrt(self.x**2+self.y**2+self.z**2)
        #self.orientation_axis=np.array([0, 0, 1])
        #self.coord_system=

    def get_savetext(self):
        text=self.art.name+"\t"+str(self.position)+"\t"+str(self.orientation_axis)+"\t"+str(self.system)+"\t"+str(self.active)+"\n"
        return text

    def read_savetext(self, text):
        info=text.split("\t")
        self.set_position(eval(info[1]))
        self.orientation_axis=eval(info[2])
        self.system=eval(info[3])
        self.active=eval(info[4])
        
    def get_r(self):
        self.r=np.sqrt(self.x**2+self.y**2+self.z**2)
   
    
    def set_position(self, newpos):
        self.position=newpos
        self.x, self.y, self.z=self.position
        self.get_r()
        
    def set_x(self, newx):
        self.x=newx
        self.position=(self.x, self.y, self.z)
        self.get_r()
        
    def set_y(self, newy):
        self.y=newy
        self.position=(self.x, self.y, self.z)
        self.get_r()
   
    def set_z(self, newz):
        self.z=newz
        self.position=(self.x, self.y, self.z)
        self.get_r()
    
    def get_distance(self, other_spin):
        dist=(self.x-other_spin.x)**2+(self.y-other_spin.y)**2+(self.z-other_spin.z)**2
        dist=np.sqrt(dist)
        return dist

    def get_relative_coordinates(self, other_spin):
        r_x=(self.x-other_spin.x)
        r_y=(self.y-other_spin.y)
        r_z=(self.z-other_spin.z)
        return [r_x, r_y, r_z]
    
    def dipolar_strength(self, other_spin):
        r12=self.get_distance(other_spin)
        omega=d*(self.art.gyro*other_spin.art.gyro/r12*self.art.spinNumber*other_spin.art.spinNumber-3)
        return omega
    
    def dipolar_matrix(self, other_spin):
        pass
    
    def add_to_lattice(self, lattice):
        pass

class LatticeType():
    def __init__(self):
        self.name=""
        self.lattice_constant = [3.45E-10, 3.45E-10, 3.45E-10]
        self.position=set()
        self.n=len(self.position)
        self.material = ''
        self.set_position_array()
        
    def save_latticetype(self, filename="libs/latticetypes.txt"):
        text=self.get_savetext()
        if os.path.exists(filename):
            f=open(filename, "r")
            lattice_list=f.read()
            f.close()
            text=lattice_list+text
        f=open(filename, "w")
        f.write(text)
        f.close()
    
    def load_latticetype(self, i, filename="libs/latticetypes.txt"):
        try:
            f=open(filename, "r")
            lattice_list=f.readlines()
            f.close()
            text=lattice_list[i]
            self.read_savetext(text)
        except:
            print "Cannot load lattice"    

    def read_savetext(self, text):
        try:
            lattice_info=text.split("\t")
            self.name=lattice_info[0]
            self.material=lattice_info[1]
            self.lattice_constant=eval(lattice_info[2])
            for elem in lattice_info[3:]:
                self.position.add(eval(elem))
            self.n=len(self.position)
            self.position_array=np.zeros((self.n, 3), dtype=float)
            self.set_position_array()
        except:
            print "Cannot load lattice"
        
    def set_position_array(self):
        self.position_array = np.zeros((self.n, 3))
        position_list=list(self.position)
        for i in range(self.n):
            x, y, z=position_list[i]
            self.position_array[i]=[x, y, z]
            
    def add_position(self, pos):
        if not list(pos) in [list(i) for i in self.position_array]:
            self.position_array = np.append(self.position_array,np.array([pos]), axis=0)
            self.n = len(self.position_array)

    def get_savetext(self):
        savetext=self.name+"\t"+self.material+"\t"+str(self.lattice_constant)
        for x in self.position: savetext += "\t"+str(x)
        savetext +="\n"
        return savetext
    

class SubLattice():
    """
    only 1 Spin type
    """
    def __init__(self, spintype, mother, latticetype, name=""):
        """
        Constructor
        """
        self.name=name
        self.mother=mother
        self.art=spintype
        self.SpinList=[]
        self.maxd=3e-9
        self.mind=0
        self.conc=0 #not in % (1 = 100%)
        self.lattice=set()
        self.n_spins=0
        self.latticetype=latticetype
        self.get_n_max()
        
    def get_n_max(self):
        self.n_max=4.0/3.0*np.pi*self.latticetype.n*(self.maxd**3-self.mind**3)/np.prod(self.latticetype.lattice_constant)

    def get_concentration(self):
        self.get_n_max()
        self.conc=self.n_spins/self.n_max
        
    def get_max_number(self, maxd, mind):
        n=np.pi*4.0/3.0*self.latticetype.n*(maxd**3-mind**3)/np.prod(self.latticetype.lattice_constant)
        return n
        
    def set_maxd(self, newmax):
        if newmax<self.maxd:
            todel = []
            for i in self.SpinList:
                if i.r > newmax:
                    todel.append(i)
            for i in todel:
                self.delete_spin(i)
            self.maxd = newmax
            self.get_concentration()
        else:
            number_new_spins=4.0/3.0*np.pi*(newmax**3-self.maxd**3)/np.prod(self.latticetype.lattice_constant)*self.latticetype.n*self.conc
            number_new_spins=int(number_new_spins+0.5)
            self.create_spins(number_new_spins, newmax, self.maxd)
            self.maxd=newmax
            self.get_concentration()
    
    def set_mind(self, newmin):
        if newmin>self.mind:
            todel = []
            for i in self.SpinList:
                if i.r<newmin:
                    todel.append(i)
            for i in todel:
                self.delete_spin(i)
            self.mind=newmin
            self.get_concentration()
        else:
            number_new_spins=4.0/4.0*np.pi*(self.mind**3-newmin**3)/np.prod(self.latticetype.lattice_constant)*self.latticetype.n*self.conc
            number_new_spins=int(number_new_spins+0.5)
            self.create_spins(number_new_spins, self.mind, newmin)
            self.mind=newmin
            self.get_concentration()
            
    def add_spin(self, spin):
        self.SpinList.append(spin)
        self.lattice.add(spin.position)
        self.n_spins += 1
        if spin.r < self.mind: self.mind = spin.r 
        if spin.r > self.maxd: self.maxd = spin.r
        self.mother.lattice.add(spin.position)
        self.mother.SpinList.append(spin)
        self.get_concentration()
            
    def delete_spin(self, spin):
        self.SpinList.remove(spin)
        self.lattice.remove(spin.position)
        self.n_spins-=1
        self.mother.SpinList.remove(spin)
        self.mother.lattice.remove(spin.position)
        self.get_concentration()
        
    def random_sites(self, mind, maxd, n):
        ind = np.random.randint(0,self.latticetype.n, n)
        lattice = np.zeros((n, 3))
        lattice[:,0] = np.random.randint(-np.ceil(maxd/self.latticetype.lattice_constant[0]), np.ceil(maxd/self.latticetype.lattice_constant[0])+1, n)
        lattice[:,1] = np.random.randint(-np.ceil(maxd/self.latticetype.lattice_constant[1]), np.ceil(maxd/self.latticetype.lattice_constant[1])+1, n)
        lattice[:,2] = np.random.randint(-np.ceil(maxd/self.latticetype.lattice_constant[2]), np.ceil(maxd/self.latticetype.lattice_constant[2])+1, n)
        lattice += self.latticetype.position_array[ind]
        lattice *= np.array(self.latticetype.lattice_constant)
        return lattice
        
    def create_spins(self, number, max_dist="maxd", min_dist="mind"):
        if max_dist=="maxd": max_dist=self.maxd
        if min_dist=="mind": min_dist=self.mind
        if number<self.get_max_number(max_dist, min_dist):
            
            lattice = self.random_sites(min_dist, max_dist, number)
            
            lattice=set(tuple(element) for element in lattice)
            lattice=lattice.difference(self.mother.lattice)
            lattice=np.array([list(t) for t in set(tuple(element) for element in lattice)])
            
            rollagain=(np.array([1]),)
            rollagain2=(np.array([1]),)
            
            while lattice.shape[0]<number or rollagain[0].shape[0]>0 or rollagain2[0].shape[0]>0:
                toadd = self.random_sites(min_dist, max_dist, number-len(lattice))
                
                lattice=np.append(lattice,toadd,axis=0)
                lattice=set(tuple(element) for element in lattice)
                lattice=lattice.difference(self.mother.lattice)
                lattice=np.array([list(t) for t in set(tuple(element) for element in lattice)])
                
                rtest=np.sqrt((lattice[:,0])**2 +(lattice[:,1])**2 +(lattice[:,2])**2)
                rollagain=np.nonzero(rtest>max_dist)
                
                toadd = self.random_sites(min_dist, max_dist, len(rollagain[0]))
                
                lattice[rollagain,:] = toadd
                
                rtest=np.sqrt((lattice[:,0])**2 +(lattice[:,1])**2 +(lattice[:,2])**2)
                rollagain2=np.nonzero(rtest<min_dist)
                
                toadd = self.random_sites(min_dist, max_dist, len(rollagain2[0]))
                
                lattice[rollagain2,:] = toadd
                
            for i in range(number):
                spin=OneSpin(self.art)
                spin.set_position(tuple(lattice[i]))
                if self.mother.lattice.isdisjoint(spin.position):
                    self.add_spin(spin)
                    
            self.get_concentration()
        else: print "Too many spins"
  
  
    def set_concentration(self, c):
        number=c*self.n_max
        number=int(number+0.5)
        while (number+1)<self.n_spins:
                ind=np.random.randint(0, self.n_spins)
                self.delete_spin(self.SpinList[ind])
        if number>self.n_spins:
            self.create_spins(number-self.n_spins)
        self.get_concentration()
        
        
class OneLattice():
    """
    superlattice
    """
    def __init__(self, material, name="Lattice"):
        self.material=material
        self.lattice=set()   
        self.SublatticeList=[]
        self.sys = phy.System()
        self.SpinList = []
        self.name=name
        
    def add_sublattice(self, spintype, lat_type, name):
        new_sublattice=SubLattice(spintype, self, lat_type, name)
        self.SublatticeList.append(new_sublattice)
        
    def create_spin(self, sublattice, number, maxd="maxd", mind="mind"):
        sublattice.create_spins(number, maxd, mind)
        
    def add_spin(self, spin, sublattice):
        if self.lattice.isdisjoint(spin.position):
            sublattice.add_spin(spin)

        
    def delete_spin(self, spin):
        for sublattice in SublatticeList:
            if spin in sublattice.SpinList:
                sublattice.delete_spin(spin)
        self.lattice-=spin.position
        
    """   
    def update_interaction_matrix(self):
        lat = np.empty((0,3))
        gyro = []
        for sublat in subDict.values():
            lat_i = []
            for spin in sublat.SpinList:
                lat_i.append(spin.position)
            gyro.extend([sublat.spintype.gyro] * len(lat_i))
            lat = np.append(lat, np.array(lat_i), axis=0)
        gyro = np.array(gyro)
        self.interactions = np.zeros((len(lat),len(lat)))
        for i in range(len(lat)-1):
            self.interactions[i,i+1:]=np.einsum('ij,ij->i',lat[i+1:]-lat[i], lat[i+1]-lat[i])**(-3)
        self.interactions += self.interactions.T
        self.interactions *= np.outer(gyro, gyro) * sconst.mu0 / (sconst.h * np.pi * 4)
    """
        
    def delete_lattice(self, sublattice):
        self.lattice -= sublattice.lattice()
        self.SublatticeList.remove(sublattice)

    def get_savetext(self):
        lattypes = []
        spintypes = []
        data="***\t "+self.name+" \t"+self.material+"\t ***\n"
        data+='-'*80+"\n"
        data+="Lattice Types: \n"
        data+='~'*80+"\n"
        for sublattice in self.SublatticeList:
            if sublattice.latticetype not in lattypes:
                lattypes.append(sublattice.latticetype)
            if sublattice.art not in spintypes:
                spintypes.append(sublattice.art)
        for latticetype in lattypes:
            data+= latticetype.get_savetext()
        data+= '-'*80 + '\n'
        data+= 'Spin Types: \n'
        data+= '~'*80 +'\n'
        for spintype in spintypes:
            data+= spintype.get_savetext()
        data+= '-'*80 +'\n'
        for sublattice in self.SublatticeList:
            #data+= '+'*80 +'\n'
            data+= '+++ Sublattice '+sublattice.name+' +++\n'
            data+= '~'*80 +'\n'
            data+= sublattice.latticetype.name + '\t' + sublattice.art.name + '\n'
            #data+= str(sublattice.maxd) + '\t' + str(sublattice.mind) + '\t' + str(sublattice.conc) + '\n'
            data+= '.'*80 +'\n'
            for spin in sublattice.SpinList:
                data += spin.get_savetext()
            data+= '+'*80 +'\n'
        data+='.'*80 +'\n'
        return data
        
    def save_lattice(self, filename="lattice0"):
        data=self.get_savetext()
        f=open(filename, 'w')
        f.write(data)
        f.close()


    def read_savetext(self, text, list_return):
        [self_info, latticetype_info, spintype_info, sublattices_info] = text.split('\n'+"-"*80+"\n")
        latticetype_info = latticetype_info.split("\n")
        spintype_info = spintype_info.split("\n")
        sublattices_info = sublattices_info.split('\n'+"+"*80+"\n")
        self.name = self_info.split("\t")[1]
        self.material = self_info.split("\t")[2]
        latticetypes = []
        spintypes = []
        for i in range(2, len(latticetype_info)):
            new_latticetype = LatticeType()
            new_latticetype.read_savetext(latticetype_info[i])
            latticetypes.append(new_latticetype)
        for i in range(2, len(spintype_info)):
            new_spintype = SpinType()
            new_spintype.read_savetext(spintype_info[i])
            spintypes.append(new_spintype)
        for sublattice_info in sublattices_info[:-1]:
            sublattice_info = sublattice_info.split("\n")
            [lt_name, st_name] = sublattice_info[2].split('\t')
            for lt in latticetypes:
                if lt.name == lt_name:
                    sublat_lt = lt
            for st in spintypes:
                if st.name == st_name:
                    sublat_st = st
            new_sublattice = SubLattice(sublat_st, self, sublat_lt)
            new_sublattice.name=sublattice_info[0].split(" ")[2]
            self.SublatticeList.append(new_sublattice)
            for spin_info in sublattice_info[4:]:
                new_spin=OneSpin(new_sublattice.art)
                new_spin.read_savetext(spin_info)
                self.add_spin(new_spin, new_sublattice)
                if new_spin.system:
                    self.sys.add_spin(new_spin)
        if list_return:
            return latticetypes, spintypes

            
            
    
    def load_lattice(self, filename, list_return = False):
        f=open(filename, "r")
        data=f.read()
        f.close()
        if list_return:
            ltli, stli = self.read_savetext(data, list_return)
            return ltli, stli
        else:
            self.read_savetext(data, list_return)
        
        
        """
        ltli = []
        stli = []
        
        latticetypeindex = data.index('Lattice Types: \n')
        spintypeindex = data.index('Spin Types: \n')
        sli = data.index('+'*80 +'\n') #?
        nlt = int((spintypeindex-latticetypeindex-2)/3)
        nst = (sli-spintypeindex-3)
        nsl = data.count('+'*80 +'\n')
        
        for i in range(nlt):
            newlt = LatticeType()
            lt_info = data[lti+2+3*i][:-1].split('\t')
            newlt.name = lt_info[0]
            newlt.lattic_constant = eval(lt_info[1])
            lt_uc = data[lti+3+3*i][:-1].split('\t')
            for pos in lt_uc:
                newlt.add_position(eval(pos))
            ltli.append(newlt)
            
        for i in range(nst):
            newst = SpinType()
            st_info = data[sti+2+i][:-1].split('\t')
            newst.name = st_info[0]
            newst.gyro = float(st_info[1])
            newst.zfs = float(st_info[2])
            newst.spinNumber = float(st_info[3])
            stli.append(newst)
        
        for i in range(nsl-1):
            start = data.index('+'*80+ '\n')
            data.remove('+'*80+ '\n')
            end = data.index('+'*80+ '\n')
            sl_name = data[start].split()[2]
            sl_info = data[start+2][:-1].split('\t')
            sl_lt = ltli[[i.name == sl_info[0] for i in ltli].index(True)]
            sl_st = stli[[i.name == sl_info[1] for i in stli].index(True)]
            self.add_sublattice(sl_st, sl_lt, sl_name)
            maxd, mind, conc = data[start+3][:-1].split('\t')
            self.subDict[sl_name].maxd = float(maxd)
            self.subDict[sl_name].mind = float(mind)
            self.subDict[sl_name].conc = float(conc)
            for j in range(start+5, end-1):
                newspin = OneSpin(sl_st)
                newspin_info = data[j].split('\t')
                newspin.set_position(eval(newspin_info[0]))
                newspin.system = newspin_info[1] == 'True'
                newspin.active = newspin_info[2] == 'True'
                newspin.orientation_axis = eval(newspin_info[3])
                self.add_spin(newspin, sl_name)
                if newspin.system:
                    self.sys.add_spin(newspin)
                
        if list_return:
            return ltli, stli
        """

class CartesianSystem():
    def __init__(self):
        self.name="default"
        self.x_axis=np.array([1, 0, 0])
        self.y_axis=np.array([0, 1, 0])
        self.z_axis=np.array([0, 0, 1])
        self.MatrixForm=np.vstack(self.x_axis, self.y_axis, self.z_axis)
    
    def save_CoordinateSystem(self, filename="libs/CoordinateSystems.txt"):
        coord_new=self.name+"\t"+str(list(self.x_axis))+"\t"+str(list(self.y_axis))+"\t"+str(list(self.z_axis))+"\n"
        if os.path.exists(filename):
            f=open(filename, "r")
            coord_list=f.read()
            f.close()
            coord_new=coord_list+coord_new
        f=open(filename, "w")
        f.write(coord_new)
        f.close()
    
    def load_CoordinateSystem(self, i, filename="libs/CoordinateSystems.txt"):
        try:
            f=open(filename, "r")
            coord_list=f.readlines()
            f.close()
            coord_info=coord_list[i]
            coord_info=coord_info.split("\t")
            if len(coord_info)==4:
                self.name=coord_info[0]
                self.x_axis=np.array(eval(coord_info[1]))
                self.y_axis=np.array(eval(coord_info[2]))
                self.z_axis=np.array(eval(coord_info[3]))
            else: print "not right format"
        except:
            print "Cannot do it"    
    
    def rotate(self, vec):
        rotated=np.dot(self.MatrixForm, vec)
        return rotated
        
    
        
